package com.atomgit.ai.proxy.service.impl;

import com.alibaba.fastjson.JSON;
import com.atomgit.ai.proxy.entity.ModelInfo;
import com.atomgit.ai.proxy.entity.UserApiKey;
import com.atomgit.ai.proxy.mapper.ModelInfoMapper;
import com.atomgit.ai.proxy.mapper.UserApiKeyMapper;
import com.atomgit.ai.proxy.param.PromptParam;
import com.atomgit.ai.proxy.processor.service.ParamAndResponseProcessor;
import com.atomgit.ai.proxy.service.ModelInfoService;
import com.atomgit.ai.proxy.util.HTTPUtil;
import com.atomgit.ai.proxy.vo.*;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpMethod;
import org.springframework.stereotype.Service;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Service
@Slf4j
@RequiredArgsConstructor
public class ModelInfoServiceImpl implements ModelInfoService {

    @Autowired
    private ModelInfoMapper modelInfoMapper;

    @Autowired
    private UserApiKeyMapper userApiKeyMapper;

    @Autowired
    private HTTPUtil httpProxyStreamUtil;

    @Override
    public List<ModelInfoVO> list() {
        List<ModelInfoVO> modelInfoVOS = new ArrayList<>();
        List<ModelInfo> modelInfos = modelInfoMapper.selectAll();
        for (ModelInfo modelInfo : modelInfos){
             ModelInfoVO modelInfoVO = new ModelInfoVO();
             BeanUtils.copyProperties(modelInfo,modelInfoVO);
             List<String> versionList = Arrays.asList(modelInfo.getVersion().split(","));
             List<ModelVersionVO> modelVersionVOList = new ArrayList<>();
             for (int i = 0; i < versionList.size(); i++) {
                String version =  versionList.get(i);
                ModelVersionVO modelVersionVO = new ModelVersionVO();
                modelVersionVO.setVersion(version);
                modelVersionVOList.add(modelVersionVO);
             }
             modelInfoVO.setVersionList(modelVersionVOList);
             modelInfoVOS.add(modelInfoVO);
        }
        return modelInfoVOS;
    }

    @Override
    public ModelAnswerVO answer(PromptParam param) throws ClassNotFoundException, InstantiationException, IllegalAccessException {
        log.info("获取指定模型回答...");
        try {
            ModelInfo modelInfo = modelInfoMapper.selectById(param.getModelId());
            // 根据userId及modelId查询指定apiKey
            UserApiKey userApiKey = userApiKeyMapper.selectByModelIdAndVersion(param.getModelId(), param.getVersion());
            // 初始化模型处理类并处理参数与返回结果
            Class<?> aClass = Class.forName(modelInfo.getProcessorClassPath());
            ParamAndResponseProcessor paramAndResponseProcessor = (ParamAndResponseProcessor) aClass.newInstance();
            String apiKey = paramAndResponseProcessor.getApiKey(modelInfo, userApiKey);
            String requestBody = paramAndResponseProcessor.generateRequestBody(param);
            String url = paramAndResponseProcessor.getUrl(modelInfo, userApiKey);
            String responseBody = httpProxyStreamUtil.sendRequest(url, HttpMethod.POST, requestBody, apiKey);
            if (responseBody.equals("Invalid API-key provided.")){
                throw new RuntimeException("Invalid API-key provided.");
            }
            return paramAndResponseProcessor.generateModelAnswerVO(responseBody);
        }catch (Exception e){
            ModelAnswerVO modelAnswerVO = new ModelAnswerVO();
            ModelAnswerOutPutVO modelAnswerOutPutVO = new ModelAnswerOutPutVO();
            modelAnswerOutPutVO.setText("目前出现了一点错误");
            modelAnswerOutPutVO.setFinish_reason("stop");
            modelAnswerVO.setOutput(modelAnswerOutPutVO);
            return modelAnswerVO;
        }
    }

    @Override
    public void answerInStream(PromptParam param, HttpServletResponse response) throws ClassNotFoundException, InstantiationException, IllegalAccessException, IOException {
        try {
            log.info("获取指定模型回答并以流式返回结果...");
            ModelInfo modelInfo = modelInfoMapper.selectById(param.getModelId());
            // 根据userId及modelId查询指定apiKey
            UserApiKey userApiKey = userApiKeyMapper.selectByModelIdAndVersion(param.getModelId(),param.getVersion());
            // 初始化模型处理类并处理参数与返回结果
            Class<?> aClass = Class.forName(modelInfo.getProcessorClassPath());
            ParamAndResponseProcessor paramAndResponseProcessor = (ParamAndResponseProcessor) aClass.newInstance();
            String apiKey = paramAndResponseProcessor.getApiKey(modelInfo,userApiKey);
            String url = paramAndResponseProcessor.getUrl(modelInfo,userApiKey);
            String requestBody = paramAndResponseProcessor.generateRequestBody(param);
            paramAndResponseProcessor.sendRequestInStream(url, HttpMethod.POST,requestBody,apiKey,response,paramAndResponseProcessor);
        }catch (Exception e){
            log.info(e.getMessage());
            ModelAnswerVO modelAnswerVO = new ModelAnswerVO();
            ModelAnswerOutPutVO modelAnswerOutPutVO = new ModelAnswerOutPutVO();
            modelAnswerOutPutVO.setText(e.getMessage());
            modelAnswerOutPutVO.setFinish_reason("stop");
            // 返回utf8格式的提示：“目前出现了一点错误”
            response.setContentType("text/event-stream;charset=UTF-8");
            modelAnswerOutPutVO.setText("目前出现了一点错误");
            modelAnswerVO.setOutput(modelAnswerOutPutVO);
            response.getWriter().println(JSON.toJSONString(modelAnswerVO));
            response.getWriter().flush();
        }
    }

    @Override
    public List<ModelInfoSelfVO> listSelf() {
        List<ModelInfoSelfVO> modelInfoSelfVOS = new ArrayList<>();
        List<UserApiKey> userApiKeys = userApiKeyMapper.selectAll();
        for (UserApiKey userApiKey : userApiKeys){
            ModelInfoSelfVO modelInfoSelfVO = new ModelInfoSelfVO();
            BeanUtils.copyProperties(userApiKey,modelInfoSelfVO);
            modelInfoSelfVO.setApiKey(userApiKey.getApiKey());
            modelInfoSelfVO.setName(userApiKey.getCustomModelName());
            modelInfoSelfVO.setModelId(userApiKey.getModelId());
            modelInfoSelfVOS.add(modelInfoSelfVO);
        }
        return modelInfoSelfVOS;
    }

}